import numpy
import SimpleITK as sitk
def command_iteration(method):
    print(f"{method.GetOptimizerIteration():3} = {method.GetMetricValue():10.5f}")
for now in [6,7,8]:#[1,2,4,5]:#[1,2,4,5]:
    #NOTE the displacmentField/transform derived is BACKWARD,
    # which goes from FIXED to MOVING   see https://discourse.itk.org/t/demons-displacement-field-backwards/1598/2 for detail
    moving=sitk.ReadImage('Masked/'+str(now)+'masked.nii',sitk.sitkFloat32)
    fixed=sitk.ReadImage('Masked/'+str(3)+'masked.nii',sitk.sitkFloat32)

    #algorithm parameter tuning
    # B-spline mesh size is the number of polynomial patches comprising the finite domain of support
    # mesh size = number of control points - spline order(3 by default)
    transformDomainMeshSize = [12] * moving.GetDimension()
    print(transformDomainMeshSize)
    tx = sitk.BSplineTransformInitializer(fixed,transformDomainMeshSize)

    print("Initial Parameters:")
    print(tx.GetParameters())

    R = sitk.ImageRegistrationMethod()
    R.SetMetricAsCorrelation()
    R.SetOptimizerAsLBFGSB(gradientConvergenceTolerance=1e-5,
                           numberOfIterations=50,
                           maximumNumberOfCorrections=5,
                           maximumNumberOfFunctionEvaluations=500,
                           costFunctionConvergenceFactor=1e+7)
    R.SetInitialTransform(tx, True)
    R.SetInterpolator(sitk.sitkLinear)

    R.AddCommand(sitk.sitkIterationEvent, lambda: command_iteration(R))

    outTx = R.Execute(fixed, moving)

    print("-------")
    print(outTx)
    print(f"Optimizer stop condition: {R.GetOptimizerStopConditionDescription()}")
    print(f" Iteration: {R.GetOptimizerIteration()}")
    print(f" Metric value: {R.GetMetricValue()}")

    sitk.WriteTransform(outTx,str(now)+'reg3'+'.mat')

    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(fixed)
    resampler.SetOutputPixelType(sitk.sitkUInt16)
    resampler.SetInterpolator(sitk.sitkLinear)
    resampler.SetDefaultPixelValue(800)
    resampler.SetTransform(outTx)

    out = resampler.Execute(moving)
    sitk.WriteImage(out,str(now)+'toward3.nii')

